import base64
import glob
import json
import logging
import os
import sys
import shutil
import subprocess
import uuid
import time
from collections import defaultdict
from threading import Lock
from flask import Flask, request, jsonify
from logging.handlers import TimedRotatingFileHandler

# 加载配置文件
CONFIG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "mcp.json")
if os.path.exists(CONFIG_FILE):
    with open(CONFIG_FILE, 'r') as f:
        config = json.load(f)
else:
    config = {
        "server": {
            "port": 8080,
            "host": "0.0.0.0",
            "max_execution_time": 30,
            "max_requests_per_minute": 10
        },
        "directories": {
            "base_dir": ".",
            "venvs_dir": "venvs",
            "temp_dir": "temp",
            "output_dir": "output",
            "logs_dir": "logs"
        },
        "security": {
            "blacklisted_modules": ["os", "subprocess", "sys", "shutil", "socket"],
            "allowed_file_extensions": [".png", ".jpg", ".jpeg", ".gif", ".svg", ".pdf", ".csv", ".json", ".txt"],
            "max_file_size": 10485760,
            "max_files_count": 10
        },
        "logging": {
            "level": "INFO",
            "format": "%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s",
            "file": "logs/mcp_server.log",
            "backup_count": 30
        },
        "dependencies": {
            "use_cached_venvs": True,
            "venv_expiry_days": 7
        }
    }

# 配置项
# 虚拟环境和临时目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
VENVS_DIR = os.path.join(BASE_DIR, config["directories"]["venvs_dir"])
TEMP_DIR = os.path.join(BASE_DIR, config["directories"]["temp_dir"])
OUTPUT_DIR = os.path.join(BASE_DIR, config["directories"]["output_dir"])
# 服务器配置
MAX_EXECUTION_TIME = config["server"]["max_execution_time"]
MAX_REQUESTS_PER_MINUTE = config["server"]["max_requests_per_minute"]
BLACKLISTED_MODULES = config["security"]["blacklisted_modules"]
SERVER_PORT = config["server"]["port"]
# 文件配置
ALLOWED_FILE_EXTENSIONS = config["security"]["allowed_file_extensions"]
MAX_FILE_SIZE = config["security"]["max_file_size"]
MAX_FILES_COUNT = config["security"]["max_files_count"]
# 日志配置
LOG_DIR = os.path.join(BASE_DIR, config["directories"]["logs_dir"])
LOG_FILE = os.path.join(BASE_DIR, config["logging"]["file"])
LOG_LEVEL = getattr(logging, config["logging"]["level"])
LOG_FORMAT = config["logging"]["format"]

# 创建日志目录（如果不存在）
os.makedirs(LOG_DIR, exist_ok=True)

# 配置日志
logger = logging.getLogger("mcp_server")
logger.setLevel(LOG_LEVEL)
# 控制台处理器
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(logging.Formatter(LOG_FORMAT))
logger.addHandler(console_handler)
# 文件处理器（按天轮转）
file_handler = TimedRotatingFileHandler(
    LOG_FILE,
    when="midnight",
    interval=1,
    backupCount=config["logging"]["backup_count"],  # 保留日志天数
    encoding="utf-8"
)
file_handler.setFormatter(logging.Formatter(LOG_FORMAT))
logger.addHandler(file_handler)

app = Flask(__name__)

# 创建必要的目录（如果不存在）
logger.info("初始化MCP服务器...")
os.makedirs(VENVS_DIR, exist_ok=True)
logger.info(f"创建虚拟环境目录: {VENVS_DIR}")
os.makedirs(TEMP_DIR, exist_ok=True)
logger.info(f"创建临时文件目录: {TEMP_DIR}")

# 简单的速率限制实现
rate_limits = defaultdict(list)
rate_limit_lock = Lock()


def is_rate_limited(ip_address):
    """
    检查IP地址是否超出速率限制

    参数:
        ip_address: 客户端IP地址

    返回:
        布尔值，表示是否被限制
    """
    with rate_limit_lock:
        current_time = time.time()
        # 移除1分钟前的请求记录
        rate_limits[ip_address] = [t for t in rate_limits[ip_address] if current_time - t < 60]
        # 检查是否达到限制
        if len(rate_limits[ip_address]) >= MAX_REQUESTS_PER_MINUTE:
            logger.warning(f"IP {ip_address} 请求频率超限")
            return True
        # 添加当前请求
        rate_limits[ip_address].append(current_time)
        return False


def is_code_safe(code):
    """
    基本安全检查，查找潜在危险的导入语句

    参数:
        code: 要检查的Python代码

    返回:
        (is_safe, reason): 元组，包含布尔值表示代码是否安全，字符串说明原因
    """
    lines = code.split('\n')
    for line in lines:
        line = line.strip()
        if line.startswith('import ') or line.startswith('from '):
            for module in BLACKLISTED_MODULES:
                if f" {module}" in line or line.startswith(f"import {module}") or line.startswith(f"from {module}"):
                    logger.warning(f"检测到危险导入: {module}")
                    return False, f"出于安全考虑，不允许导入{module}模块"
    logger.debug("代码安全检查通过")
    return True, ""


def find_and_encode_files(temp_path):
    """
    查找并编码用户代码生成的文件

    参数:
        temp_path: 临时目录路径，包含用户代码和生成的文件

    返回:
        包含文件信息的字典列表
    """
    files_info = []

    # 查找所有符合允许扩展名的文件
    for ext in ALLOWED_FILE_EXTENSIONS:
        pattern = os.path.join(temp_path, f"*{ext}")
        matching_files = glob.glob(pattern)

        for file_path in matching_files:
            # 跳过代码文件本身
            if os.path.basename(file_path) == "code.py":
                continue

            # 检查文件大小
            file_size = os.path.getsize(file_path)
            if file_size > MAX_FILE_SIZE:
                logger.warning(f"文件超过大小限制: {file_path} ({file_size} 字节)")
                continue

            # 如果文件总数达到限制，停止添加
            if len(files_info) >= MAX_FILES_COUNT:
                logger.warning(f"达到最大文件数量限制 ({MAX_FILES_COUNT})，忽略其余文件")
                break

            try:
                # 读取并编码文件内容
                with open(file_path, 'rb') as f:
                    file_content = f.read()

                # 确定文件的MIME类型
                ext = os.path.splitext(file_path)[1].lower()
                mime_type = {
                    '.png': 'image/png',
                    '.jpg': 'image/jpeg',
                    '.jpeg': 'image/jpeg',
                    '.gif': 'image/gif',
                    '.svg': 'image/svg+xml',
                    '.pdf': 'application/pdf',
                    '.csv': 'text/csv',
                    '.json': 'application/json',
                    '.txt': 'text/plain'
                }.get(ext, 'application/octet-stream')

                # 添加文件信息
                files_info.append({
                    'name': os.path.basename(file_path),
                    'size': file_size,
                    'mime_type': mime_type,
                    'content': base64.b64encode(file_content).decode('utf-8')
                })

                logger.info(f"找到并编码文件: {os.path.basename(file_path)} ({file_size} 字节)")

            except Exception as e:
                logger.error(f"处理文件时出错 {file_path}: {str(e)}")

    return files_info


@app.route('/run', methods=['POST', "GET"])
def run_code():
    """
    处理代码执行请求的主要API端点
    """
    try:
        # 获取客户端IP用于速率限制
        client_ip = request.remote_addr
        logger.info(f"收到来自 {client_ip} 的代码执行请求")

        # 检查速率限制
        if is_rate_limited(client_ip):
            logger.warning(f"IP {client_ip} 请求频率超限")
            return jsonify({"error": "请求频率超限，请稍后再试"}), 429

        data = request.json
        if not data:
            logger.error("请求中未提供数据")
            return jsonify({"error": "未提供数据"}), 400

        # 提取代码和依赖
        code = data.get('code')
        dependencies = data.get('dependencies', [])
        logger.info(f"依赖列表: {dependencies}")

        if not code:
            logger.error("请求中未提供代码")
            return jsonify({"error": "未提供代码"}), 400

        # 检查代码安全性
        is_safe, reason = is_code_safe(code)
        if not is_safe:
            logger.warning(f"代码安全检查失败: {reason}")
            return jsonify({"error": reason}), 403

        # 为本次运行创建唯一ID
        run_id = str(uuid.uuid4())
        logger.info(f"为请求创建唯一ID: {run_id}")
        venv_path = os.path.join(VENVS_DIR, run_id)
        temp_path = os.path.join(TEMP_DIR, run_id)
        os.makedirs(temp_path, exist_ok=True)
        logger.info(f"创建临时目录: {temp_path}")

        # 将代码保存到临时文件
        code_file_path = os.path.join(temp_path, "code.py")
        with open(code_file_path, "w") as f:
            f.write(code)
        logger.info(f"代码已保存到: {code_file_path}")

        # 创建虚拟环境
        logger.info(f"开始创建虚拟环境: {venv_path}...")
        result = create_venv(venv_path)
        if result.get('error'):
            logger.error(f"创建虚拟环境失败: {result.get('error')}")
            cleanup(venv_path, temp_path)
            return jsonify(result), 500
        logger.info("虚拟环境创建成功")

        # 安装依赖
        if dependencies:
            logger.info(f"开始安装依赖: {dependencies}")
            result = install_dependencies(venv_path, dependencies)
            if result.get('error'):
                logger.error(f"安装依赖失败: {result.get('error')}")
                cleanup(venv_path, temp_path)
                return jsonify(result), 500
            logger.info("依赖安装成功")

        # 运行代码
        logger.info("开始执行代码...")
        result = execute_code(venv_path, code_file_path, temp_path)
        logger.info(f"代码执行完成，状态码: {result.get('exit_code', '未知')}")

        # 查找并处理生成的文件
        if result.get('exit_code') == 0:
            logger.info("查找代码生成的文件...")
            files = find_and_encode_files(temp_path)
            if files:
                logger.info(f"找到 {len(files)} 个文件")
                result['files'] = files

        # 清理环境
        logger.info("开始清理环境...")
        cleanup(venv_path, temp_path)
        logger.info("环境清理完成")

        return jsonify(result)

    except Exception as e:
        logger.exception(f"处理请求时发生错误: {str(e)}")
        return jsonify({"error": str(e)}), 500


def create_venv(venv_path):
    """
    创建虚拟环境

    参数:
        venv_path: 虚拟环境路径

    返回:
        包含操作结果的字典
    """
    try:
        logger.debug(f"运行命令创建虚拟环境: python -m virtualenv {venv_path}")
        subprocess.run(
            ["python", "-m", "virtualenv", venv_path],
            check=True,
            capture_output=True,
            text=True
        )
        logger.debug("虚拟环境创建命令执行成功")
        return {"success": "虚拟环境创建成功"}
    except subprocess.CalledProcessError as e:
        logger.error(f"创建虚拟环境过程中出错: {e.stderr}")
        return {"error": f"创建虚拟环境失败: {e.stderr}"}


def install_dependencies(venv_path, dependencies):
    """
    在虚拟环境中安装依赖

    参数:
        venv_path: 虚拟环境路径
        dependencies: 要安装的依赖列表

    返回:
        包含操作结果的字典
    """
    try:
        # 根据操作系统确定pip可执行文件路径
        pip_path = os.path.join(venv_path, "bin", "pip") if not os.name == 'nt' else os.path.join(venv_path, "Scripts", "pip")
        logger.debug(f"确定pip路径: {pip_path}")

        # 安装每个依赖
        for dep in dependencies:
            # 基本验证包名以防止命令注入
            if not all(c.isalnum() or c in '._-=<>' for c in dep):
                logger.warning(f"检测到无效的包名: {dep}")
                return {"error": f"无效的包名: {dep}"}

            logger.info(f"安装依赖: {dep}")
            result = subprocess.run(
                [pip_path, "install", dep],
                check=True,
                capture_output=True,
                text=True
            )
            logger.debug(f"依赖 {dep} 安装输出: {result.stdout}")

        return {"success": "依赖安装成功"}
    except subprocess.CalledProcessError as e:
        logger.error(f"安装依赖时出错: {e.stderr}")
        return {"error": f"安装依赖失败: {e.stderr}"}


def execute_code(venv_path, code_file_path, work_dir):
    """
    在虚拟环境中执行代码

    参数:
        venv_path: 虚拟环境路径
        code_file_path: 代码文件路径
        work_dir: 工作目录，代码将在此目录下执行，文件将生成在此目录中

    返回:
        包含执行结果的字典
    """
    try:
        # 根据操作系统确定python可执行文件路径
        python_path = os.path.join(venv_path, "bin", "python") if not os.name == 'nt' else os.path.join(venv_path, "Scripts", "python")
        logger.debug(f"确定Python路径: {python_path}")

        # 执行代码
        logger.info(f"开始执行代码文件: {code_file_path}")
        logger.debug(f"命令: {python_path} {code_file_path}")
        # 设置工作目录，这样代码生成的文件会保存在临时目录中
        result = subprocess.run(
            [python_path, code_file_path],
            check=True,
            capture_output=True,
            text=True,
            timeout=MAX_EXECUTION_TIME,  # 设置代码执行超时
            cwd=work_dir  # 设置工作目录
        )

        logger.info(f"代码执行成功，退出码: {result.returncode}")
        if result.stdout:
            logger.debug(f"标准输出: {result.stdout[:200]}...")
        if result.stderr:
            logger.debug(f"标准错误: {result.stderr[:200]}...")

        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "exit_code": result.returncode
        }
    except subprocess.TimeoutExpired:
        logger.warning(f"代码执行超时（{MAX_EXECUTION_TIME}秒）")
        return {"error": f"代码执行超时（{MAX_EXECUTION_TIME}秒）"}
    except subprocess.CalledProcessError as e:
        logger.error(f"代码执行失败，退出码: {e.returncode}")
        if e.stdout:
            logger.debug(f"标准输出: {e.stdout[:200]}...")
        if e.stderr:
            logger.debug(f"标准错误: {e.stderr[:200]}...")

        return {
            "stdout": e.stdout,
            "stderr": e.stderr,
            "exit_code": e.returncode,
            "error": "代码执行失败"
        }


def cleanup(venv_path, temp_path):
    """
    清理虚拟环境和临时文件

    参数:
        venv_path: 虚拟环境路径
        temp_path: 临时文件路径
    """
    try:
        if os.path.exists(venv_path):
            logger.debug(f"删除虚拟环境目录: {venv_path}")
            shutil.rmtree(venv_path)
        if os.path.exists(temp_path):
            logger.debug(f"删除临时文件目录: {temp_path}")
            shutil.rmtree(temp_path)
    except Exception as e:
        logger.error(f"清理失败: {str(e)}")
        print(f"警告: 清理失败: {str(e)}")


@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    logger.debug("收到健康检查请求")
    return jsonify({"status": "healthy"})


@app.route('/logs', methods=['GET'])
def view_logs():
    """查看最新日志的API端点"""
    try:
        # 获取参数
        lines = request.args.get('lines', default=100, type=int)
        level = request.args.get('level', default='INFO').upper()

        # 验证和限制lines参数
        if lines <= 0:
            lines = 100
        elif lines > 1000:
            lines = 1000

        # 读取日志文件
        if os.path.exists(LOG_FILE):
            with open(LOG_FILE, 'r', encoding='utf-8') as f:
                log_content = f.readlines()

            # 筛选日志级别
            if level in ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']:
                log_content = [line for line in log_content if f" - {level}: " in line]

            # 获取最新的N行
            log_content = log_content[-lines:]

            return jsonify({
                "logs": log_content,
                "count": len(log_content)
            })
        else:
            return jsonify({"error": "日志文件不存在"}), 404
    except Exception as e:
        logger.exception(f"获取日志时发生错误: {str(e)}")
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    logger.info("MCP Server starting...")
    # 确保目录存在
    os.makedirs(VENVS_DIR, exist_ok=True)
    os.makedirs(TEMP_DIR, exist_ok=True)
    app.run(host=config["server"]["host"], port=SERVER_PORT, debug=False)
